#!/usr/bin/python2

import subprocess
import re
from os import path

SITES = dict(
    google = 'www.google.com',
    duckduckgo = 'duckduckgo.com',
    github = 'github.com',
    wikipedia = 'wikipedia.org',
    arstechnica = 'arstechnica.com',
    reddit = 'reddit.com',
    hn = 'news.ycombinator.com',
    servo = 'servo.org',
    rustlang = 'www.rust-lang.org',
    wapo = 'www.washingtonpost.com',
    twitter = 'twitter.com',
    stackoverflow = 'stackoverflow.com',
)

def extract_certs(line):
    if 'Server cert is [' not in line:
        return None

    chain = []
    for certm in re.finditer('Certificate\(\[([\d, ]+)\]\)', line):
        bytes = certm.group(1).split(', ')
        bytes = map(int, bytes)
        cert = ''.join(map(chr, bytes))
        chain.append(cert)
    return chain

def collect(hostname):
    subp = subprocess.Popen([
        './target/debug/examples/tlsclient',
        '--verbose',
        '--http',
        hostname],
        stderr = subprocess.PIPE,
        stdout = subprocess.PIPE)
    stdout, stderr = subp.communicate()

    certs = None

    for line in stderr.splitlines():
        found = extract_certs(line)
        if found:
            certs = found

    return certs

if __name__ == '__main__':
    certfile = lambda name, i: 'src/testdata/cert-%s.%d.der' % (name, i)

    for name, hostname in SITES.items():
        if path.exists(certfile(name, 0)):
            continue
        certchain = collect(hostname)

        for i, cert in enumerate(certchain):
            open(certfile(name, i), 'w').write(cert)

        print '#[test]'
        print 'fn test_%s_cert() {' % name

        for i in range(len(certchain)):
            print '    let cert%d = key::Certificate(include_bytes!("testdata/cert-%s.%d.der").to_vec());' % (i, name, i)
        print '    let chain = [ %s ];' % (', '.join('cert%d' % i for i in range(len(certchain))))
        print '    let mut anchors = anchors::RootCertStore::empty();'
        print '    anchors.add_trust_anchors(&webpki_roots::ROOTS);'
        print '    bench(100, "verify_server_cert(%s)", ' % name
        print '          || (),'
        print '          |_| verify::verify_server_cert(&anchors, &chain[..], "%s", &[]).unwrap());' % hostname
        print '}'
        print
